import warnings

warnings.filterwarnings('ignore')
import numpy as np
import tensorflow as tf
import gym
from gym.spaces import Box
from gym.wrappers.monitor import Monitor
# from scipy.misc import imresize
import random
import cv2
import time
import logging
import os
import sys

# We
# define
# the
# Class
# EnvWrapper and define
# some
# of
# the
# environment
# wrapper
# functions
#
# In[2]:


class EnvWrapper:

    # First we define the __init__ method and initialize variables

    def __init__(self, env_name, debug=False):

        # environment name
        self.env_name = env_name

        # initialize the gym environment
        self.env = gym.make(env_name)

        # get the action space
        self.action_space = self.env.action_space

        # get the observation space
        self.observation_space = Box(low=0, high=255, shape=(84, 84, 4))

        # initialize frame_num for storing the frame count
        self.frame_num = 0

        # For recording the game screen
        self.monitor =  Monitor(self.env,"./test/",force=True)

        # initialize frames
        self.frames = np.zeros((84, 84, 4), dtype=np.uint8)

        # initialize a boolean called debug when set true last few frames will be displayed
        self.debug = debug

        if self.debug:
            cv2.startWindowThread()
            cv2.namedWindow("Game")

    # we define the function called step where we perform some action in the
    # environment, receive reward and move to the next state
    # step function will take the current state as input and returns the preprocessed frame as next state

    def step(self, a):
        ob, reward, done, xx = self.env.step(a)
        return self.process_frame(ob), reward, done, xx

    # We define the helper function called reset for resetting the environment
    # after resetting it will return the preprocessed game screen

    def reset(self):
        self.frame_num = 0
        return self.process_frame(self.env.reset())

    # next we define another helper function for rendering the environment
    def render(self):
        return self.env.render()

    # now we define the function called process_frame for preprocessing the frame

    def process_frame(self, frame):

        # convert the image to gray
        state_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # change the size
        state_resized = cv2.resize(state_gray, (84, 110))

        # resize
        gray_final = state_resized[16:100, :]

        if self.frame_num == 0:
            self.frames[:, :, 0] = gray_final
            self.frames[:, :, 1] = gray_final
            self.frames[:, :, 2] = gray_final
            self.frames[:, :, 3] = gray_final

        else:
            self.frames[:, :, 3] = self.frames[:, :, 2]
            self.frames[:, :, 2] = self.frames[:, :, 1]
            self.frames[:, :, 1] = self.frames[:, :, 0]
            self.frames[:, :, 0] = gray_final

        # increment the frame_num counter

        self.frame_num += 1

        if self.debug:
            cv2.imshow('Game', gray_final)

        return self.frames.copy()